#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
自定义分类训练器对话框
"""

import sys
import os
import json
from PyQt5.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QFormLayout, 
                             QLineEdit, QComboBox, QPushButton, QLabel, 
                             QTextEdit, QMessageBox, QFileDialog, QTableWidget,
                             QTableWidgetItem, QHeaderView, QAbstractItemView,
                             QApplication)
from PyQt5.QtCore import Qt
from database import Database


class CustomClassifierDialog(QDialog):
    """自定义分类训练器对话框"""
    
    def __init__(self, parent=None, db=None):
        super().__init__(parent)
        self.db = db or Database()
        self.categories = []  # 一级分类列表
        self.subcategories = {}  # 二级分类字典 {一级分类: [二级分类列表]}
        self.setup_ui()
        self.init_categories()
        
    def setup_ui(self):
        """设置UI界面"""
        layout = QVBoxLayout()
        
        # 标题
        title_label = QLabel("自定义分类训练器")
        title_label.setStyleSheet("font-size: 18px; font-weight: bold;")
        layout.addWidget(title_label)
        
        # 说明文本
        info_label = QLabel("通过提供文本描述和对应的分类标签来训练自定义分类模型。\n模型将学习根据文本内容自动分类交易记录。")
        info_label.setWordWrap(True)
        layout.addWidget(info_label)
        
        # 数据表格
        self.data_table = QTableWidget(0, 4)
        self.data_table.setHorizontalHeaderLabels(["文本描述", "一级分类", "二级分类", "操作"])
        self.data_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.Stretch)
        self.data_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeToContents)
        self.data_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeToContents)
        self.data_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeToContents)
        self.data_table.verticalHeader().setVisible(False)
        self.data_table.setSelectionBehavior(QAbstractItemView.SelectRows)
        layout.addWidget(self.data_table)
        
        # 按钮布局
        button_layout = QHBoxLayout()
        
        # 添加行按钮
        self.add_row_button = QPushButton("添加行")
        self.add_row_button.clicked.connect(self.add_row)
        button_layout.addWidget(self.add_row_button)
        
        # 从数据库加载按钮
        self.load_from_db_button = QPushButton("从数据库加载样本")
        self.load_from_db_button.clicked.connect(self.load_samples_from_db)
        button_layout.addWidget(self.load_from_db_button)
        
        # 导入按钮
        self.import_button = QPushButton("导入样本文件")
        self.import_button.clicked.connect(self.import_samples)
        button_layout.addWidget(self.import_button)
        
        # 导出按钮
        self.export_button = QPushButton("导出样本文件")
        self.export_button.clicked.connect(self.export_samples)
        button_layout.addWidget(self.export_button)
        
        button_layout.addStretch()
        layout.addLayout(button_layout)
        
        # 训练按钮
        self.train_button = QPushButton("训练模型")
        self.train_button.clicked.connect(self.train_model)
        layout.addWidget(self.train_button)
        
        self.setLayout(layout)
        
        # 初始化表格
        self.init_table()
        
    def init_table(self):
        """初始化表格"""
        # 添加几行空数据
        for i in range(5):
            self.add_row()
            
    def add_row(self):
        """添加新行"""
        row_count = self.data_table.rowCount()
        self.data_table.insertRow(row_count)
        
        # 创建文本输入框
        text_edit = QLineEdit()
        self.data_table.setCellWidget(row_count, 0, text_edit)
        
        # 创建分类下拉框
        category_combo = QComboBox()
        category_combo.addItems(self.categories)
        
        # 创建二级分类下拉框
        subcategory_combo = QComboBox()
        # 设置默认的二级分类选项
        if self.categories:
            subcategory_combo.addItems(self.subcategories.get(self.categories[0], []))
        
        # 创建删除按钮
        delete_btn = QPushButton("删除")
        delete_btn.clicked.connect(lambda _, r=row_count: self.delete_row(r))
        
        # 将控件添加到表格
        self.data_table.setCellWidget(row_count, 1, category_combo)
        self.data_table.setCellWidget(row_count, 2, subcategory_combo)
        self.data_table.setCellWidget(row_count, 3, delete_btn)
        
        # 连接一级分类变化信号，更新二级分类选项
        category_combo.currentTextChanged.connect(
            lambda text: self.update_subcategories(subcategory_combo, text)
        )
        
    def delete_row(self, row):
        """删除指定行"""
        self.data_table.removeRow(row)
            
    def load_samples_from_db(self):
        """从数据库加载样本数据"""
        try:
            transactions = self.db.get_transactions()
            recent_transactions = transactions[-50:] if len(transactions) > 50 else transactions
            self.data_table.setRowCount(0)
            
            for trans in recent_transactions:
                # 提取交易描述（收款人+备注）
                description = f"{trans[5]} {trans[7]}".strip()  # payee + note
                if description and trans[4]:
                    row_count = self.data_table.rowCount()
                    self.data_table.insertRow(row_count)
                    
                    # 创建控件
                    description_edit = QLineEdit(description)
                    category_combo = QComboBox()
                    category_combo.addItems(self.categories)
                    category_combo.setCurrentText(trans[4])
                    
                    subcategory_combo = QComboBox()
                    # 获取当前一级分类的二级分类选项
                    subcategory_options = self.subcategories.get(trans[4], [])
                    subcategory_combo.addItems(subcategory_options)
                    
                    delete_btn = QPushButton("删除")
                    delete_btn.clicked.connect(lambda _, r=row_count: self.delete_row(r))
                    
                    # 将控件添加到表格
                    self.data_table.setCellWidget(row_count, 0, description_edit)
                    self.data_table.setCellWidget(row_count, 1, category_combo)
                    self.data_table.setCellWidget(row_count, 2, subcategory_combo)
                    self.data_table.setCellWidget(row_count, 3, delete_btn)
                    
                    # 连接一级分类变化信号，更新二级分类选项
                    category_combo.currentTextChanged.connect(
                        lambda text: self.update_subcategories(subcategory_combo, text)
                    )
            
            QMessageBox.information(self, "成功", f"已加载 {len(recent_transactions)} 条交易记录")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"加载数据失败: {str(e)}")
            
    def import_samples(self):
        """导入样本文件"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, '选择样本文件', '', 'JSON files (*.json)')
            
        if file_path:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    samples = json.load(f)
                
                # 清空现有数据
                self.data_table.setRowCount(0)
                
                # 添加样本到表格
                for sample in samples:
                    row_count = self.data_table.rowCount()
                    self.data_table.insertRow(row_count)
                    
                    # 创建分类下拉框
                    category_combo = QComboBox()
                    categories = [cat[1] for cat in self.db.get_categories()]
                    category_combo.addItems(categories)
                    category_combo.setCurrentText(sample.get('category', '餐饮'))
                    
                    # 创建二级分类下拉框
                    subcategory_combo = QComboBox()
                    subcategory_combo.addItems(["外卖", "堂食", "零食", "其他"])
                    subcategory_combo.setCurrentText(sample.get('subcategory', '其他'))
                    
                    # 将控件添加到表格
                    self.data_table.setCellWidget(row_count, 0, QLineEdit(sample.get('text', '')))
                    self.data_table.setCellWidget(row_count, 1, category_combo)
                    self.data_table.setCellWidget(row_count, 2, subcategory_combo)
                    
                QMessageBox.information(self, "成功", f"已导入 {len(samples)} 条样本数据")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"导入样本文件时出错: {str(e)}")
                
    def export_samples(self):
        """导出样本文件"""
        file_path, _ = QFileDialog.getSaveFileName(
            self, '保存样本文件', '', 'JSON files (*.json)')
            
        if file_path:
            try:
                # 收集表格数据
                samples = []
                for row in range(self.data_table.rowCount()):
                    # 获取文本描述
                    text_widget = self.data_table.cellWidget(row, 0)
                    text = text_widget.text() if text_widget else ""
                    
                    # 获取一级分类
                    category_widget = self.data_table.cellWidget(row, 1)
                    category = category_widget.currentText() if category_widget else "餐饮"
                    
                    # 获取二级分类
                    subcategory_widget = self.data_table.cellWidget(row, 2)
                    subcategory = subcategory_widget.currentText() if subcategory_widget else "其他"
                    
                    if text:  # 只导出有文本的行
                        samples.append({
                            'text': text,
                            'category': category,
                            'subcategory': subcategory
                        })
                
                # 保存到文件
                with open(file_path, 'w', encoding='utf-8') as f:
                    json.dump(samples, f, ensure_ascii=False, indent=2)
                
                QMessageBox.information(self, "成功", f"已导出 {len(samples)} 条样本数据到 {file_path}")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"导出样本文件时出错: {str(e)}")
                
    def train_model(self):
        """训练模型"""
        try:
            # 收集训练数据
            training_data = []
            for row in range(self.data_table.rowCount()):
                # 获取文本描述
                text_widget = self.data_table.cellWidget(row, 0)
                text = text_widget.text() if text_widget else ""
                
                # 获取一级分类
                category_widget = self.data_table.cellWidget(row, 1)
                category = category_widget.currentText() if category_widget else "餐饮"
                
                # 获取二级分类
                subcategory_widget = self.data_table.cellWidget(row, 2)
                subcategory = subcategory_widget.currentText() if subcategory_widget else "其他"
                
                if text:  # 只使用有文本的行
                    training_data.append((text, category, subcategory))
            
            if len(training_data) < 5:
                QMessageBox.warning(self, "警告", "至少需要5个样本才能训练模型")
                return
            
            # 这里应该调用实际的训练逻辑
            # 由于训练逻辑比较复杂，我们在这里只做一个模拟
            QMessageBox.information(self, "训练完成", f"模型训练完成！\n使用了 {len(training_data)} 个样本。\n\n在实际应用中，这里会调用自定义分类训练器来训练BERT模型。")
            
            # 关闭对话框
            self.accept()
        except Exception as e:
            QMessageBox.critical(self, "错误", f"训练模型时出错: {str(e)}")
    
    def init_categories(self):
         """初始化分类数据"""
         # 获取所有一级分类
         self.categories = [cat[1] for cat in self.db.get_categories()]
         
         # 获取所有二级分类，构建字典
         self.subcategories = {}
         for cat_name in self.categories:
             subcats = [subcat[2] for subcat in self.db.get_subcategories(cat_name)]
             self.subcategories[cat_name] = subcats
     
    def update_subcategories(self, subcategory_combo, category):
         """更新二级分类下拉框选项"""
         subcategory_combo.clear()
         if category in self.subcategories:
             subcategory_combo.addItems(self.subcategories[category])